package org.finesys.chat.core.base.factory;

import lombok.extern.slf4j.Slf4j;
import org.finesys.ai.constants.ModelProviderEnum;
import org.finesys.ai.entity.AigcModel;
import org.finesys.chat.core.base.client.*;
import org.finesys.chat.core.base.mapper.ModelConverter;
import org.finesys.common.core.exception.BusinessException;
import org.finesys.common.des.uitls.SM4Util;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 模型管理工厂类
 */
@Slf4j
@Component
public class ChatModelFactory {
    // 模型对应客户端集合
    private final Map<String, ChatClient> chatClientMap = new ConcurrentHashMap<>();
    //模型服务集合
    private final Map<String, AigcModel> modelMap = new ConcurrentHashMap<>();
    //加密密钥初始化向量
    @Value("${chat.model.sm4.iv:0123456789abcdeffedcba9876543210}")
    String iv;

    /**
     * 日志请求
     */
    @Value("${chat.model.log-requests:true}")
    boolean logRequests;

    /**
     * 日志响应
     */
    @Value("${chat.model.log-responses:false}")
    boolean logResponses;

    public ChatModelFactory() {
    }

    /**
     * 初始化方法
     */
    public void initModel(List<AigcModel> modelList) {
        if (!modelMap.isEmpty()) {
            modelMap.clear();
        }
        modelList.forEach(this::addOrUpdate);
    }


    public void addOrUpdate(AigcModel aigcModel) {
        String id = aigcModel.getId();
        modelMap.put(id, aigcModel);
        //转换对象
        ChatModel chatModel = ModelConverter.INSTANCE.convert(aigcModel);
        //注册客户端信息
        if (chatClientMap.containsKey(id)) {
            unregisterChatClient(chatModel);
            return;
        }
        chatModel.setLogRequests(logRequests);
        chatModel.setLogResponses(logResponses);
        // 创建模型服务
        registerChatClient(chatModel);
    }

    /**
     * 解密
     */
    private void decrypt(ChatModel chatModel) {
        String apiKey = chatModel.getApiKey();
        if (apiKey != null && !"".equals(apiKey)) {
            String apiSecret = chatModel.getApiSecret();
            try {
                chatModel.setApiKey(SM4Util.decryptCbc(apiSecret, iv, apiKey));
            } catch (Exception e) {
                throw new BusinessException("解密失败");
            }
        }
    }


    /**
     * 取消注册聊天客户端
     *
     * @param aigcModel 聊天客户端的模型
     */
    public void unregisterAigcModel(AigcModel aigcModel) {
        String modelId = aigcModel.getId();
        if (modelMap.containsKey(modelId)) {
            modelMap.remove(modelId);
        }
        if (chatClientMap.containsKey(modelId)) {
            chatClientMap.get(modelId).unregister();
            chatClientMap.remove(modelId);
        }
    }

    public AigcModel getModel(String id) {
        return modelMap.get(id);
    }

    public boolean containsModel(String id) {
        return modelMap.containsKey(id);
    }

    // 初始化模型服务
    public void init(List<ChatClient> chatClients) {
        log.info("初始化模型服务");
        chatClients.forEach(client -> {
            chatClientMap.put(client.getKey(), client);
        });
    }

    /**
     * 注册聊天客户端
     *
     * @param chatModels 聊天客户端模型列表
     */
    public void registerChatClients(List<ChatModel> chatModels) {
        log.info("注册模型服务");
        chatModels.forEach(this::registerChatClient);
    }

    /**
     * 判断是否存在
     */
    public void registerChatClient(ChatModel chatModel) {
        String modelId = chatModel.getModelId();
        if (chatClientMap.containsKey(modelId)) {
            chatClientMap.get(modelId);
            return;
        }
        try {
            //解密密钥
            decrypt(chatModel);
            // 创建模型服务
            ChatClient chatClient = createChatClient(chatModel);
            chatClientMap.put(modelId, chatClient);
        } catch (Exception e) {
            log.error("注册模型服务失败", e);
        }
    }

    /**
     * 取消注册聊天客户端
     *
     * @param chatModel 聊天客户端的模型
     */
    public void unregisterChatClient(ChatModel chatModel) {
        String modelId = chatModel.getModelId();
        if (chatClientMap.containsKey(modelId)) {
            chatClientMap.get(modelId).unregister();
            chatClientMap.remove(modelId);
        }
    }

    /**
     * 更新注册聊天客户端
     */

    public void updateChatClient(ChatModel chatModel) {
        String modelId = chatModel.getModelId();
        if (!chatClientMap.containsKey(modelId)) {
            throw new IllegalArgumentException("模型服务不存在:" + chatModel.getModelId());
        }
        chatClientMap.get(modelId).updateKey(chatModel);
        chatClientMap.put(modelId, createChatClient(chatModel));
    }

    /**
     * 根据给定的key获取ChatClient对象
     *
     * @param key ChatClient对象的唯一标识符
     * @return 返回对应的ChatClient对象
     * @throws IllegalArgumentException 如果给定的key对应的ChatClient对象不存在，则抛出IllegalArgumentException异常
     */
    public ChatClient getChatClient(String key) {
        ChatClient chatClient = chatClientMap.get(key);
        if (chatClient == null) {
            throw new IllegalArgumentException("模型服务不存在:" + key);
        }
        return chatClient;
    }

    /**
     * 创建会话客户端对象
     */

    public ChatClient createChatClient(ChatModel chatModel) {
        String provider = chatModel.getProvider();
        ModelProviderEnum providerEnum = ModelProviderEnum.getEnum(provider);
        try {
            switch (providerEnum) {
                case DEEPSEEK:
                case OLLAMA:
                    return new DeepSeekChatClient(chatModel);
                case ZHI_PU:
                    return new ZhipuChatClient(chatModel);
                case DOU_BAO:
                    return new DoubaoChatClient(chatModel);
                default:
                    return new OpenAiChatClient(chatModel);
            }
        } catch (Exception e) {
            throw new BusinessException(String.format("%s模型服务创建失败:%s", provider, e.getMessage()));
        }
    }
}
